import json

import tornado.websocket
from loguru import logger

from websocket.msghandler import MsgContext, MessageType
from websocket.msghandler.initHandler import InitHandler
from websocket.msghandler.multiTaskIndicators import MultiTaskIndicators
from websocket.msghandler.multipleDrawingHandler import MultipleDrawingHandler
from websocket.msghandler.multipleExcelHandler import MultipleExcelHandler
from websocket.msghandler.multipleTaskHandler import TaskHandler
from websocket.msghandler.simpleTaskHandler import SimpleTaskHandler
from websocket.msghandler.simpleTaskIndicators import SimpleTaskIndicators
from websocket.msghandler.statisticalHandler import StatisticalHandler


# 将类对象转换为字符串
def class_to_string(obj):
    return f"{obj.__module__}.{obj.__name__}"


class WebSocketHandler(tornado.websocket.WebSocketHandler):
    def __init__(self, application, request):
        super().__init__(application, request)
        MsgContext.callback = self.sendMsg
        self.handlers = {
            MessageType.INIT_RENDER_MSG.value: InitHandler,
            MessageType.MULTIPLE_TASK_EXECUTE.value: TaskHandler,
            MessageType.MULTI_TASK_INDICATORS.value: MultiTaskIndicators,
            MessageType.SIMPLE_TASK_RECEIVE.value: SimpleTaskHandler,
            MessageType.SIMPLE_TASK_INDICATORS.value: SimpleTaskIndicators,
            MessageType.MULTI_STATISTICAL_CALC.value: StatisticalHandler,
            MessageType.MULTI_DATA_2_EXCEL.value: MultipleExcelHandler,
            MessageType.MULTI_Drawing.value: MultipleDrawingHandler
        }

    def open(self):
        self.write_message(self.handlers[MessageType.INIT_RENDER_MSG.value].handle())

    def sendMsg(self, msg):
        self.write_message(json.dumps(msg))

    def on_message(self, message):
        message = json.loads(message)
        message_type = message['messageType']
        self.handlers[message_type].handle(message)

    def on_close(self):
        logger.info("WebSocket closed")
        # socket断开时，清空单个任务保存的数据
        MsgContext.single_populations = []
        MsgContext.simple_problem_objs = []

    def check_origin(self, origin: str) -> bool:
        return True
